# capture_frames.py

import cv2
import os
import time

DURA = 30  # seconds
SAMPLES_FPS = 10  # 每X秒抽取1帧

def extract_key_frames(camera_id,rtsp_url, duration=DURA):
    """
    从视频流中抽取三帧关键帧，每10秒一帧
    """
    
    t = int(time.time())
    temp_dir = f"/data/tmp/analysis_{camera_id}"
    os.makedirs(temp_dir, exist_ok=True)

    video_file = os.path.join(temp_dir, f"{t}_detection.mp4")
    
    print(f"📡正在请求视频流 {rtsp_url}")
     
    cap = cv2.VideoCapture(rtsp_url,cv2.CAP_FFMPEG)
    cap.set(cv2.CAP_PROP_OPEN_TIMEOUT_MSEC, 10000)   # 连接超时 10s
    cap.set(cv2.CAP_PROP_READ_TIMEOUT_MSEC, 10000)   # 读取超时 10s

    if not cap.isOpened():
        raise IOError(f"无法打开视频流: {rtsp_url}")
    
    fps = cap.get(cv2.CAP_PROP_FPS) or 25
    
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    # fourcc = cv2.VideoWriter_fourcc(*'H264')

    out = cv2.VideoWriter(video_file, fourcc, fps, (width, height))
    if not out.isOpened():
        cap.release()
        raise IOError(f"❌ 无法创建视频写入器: {video_file}")
    
    print(f"📹 视频 FPS: {fps}")
    
    key_frames = []

    frame_idx = 0
    while frame_idx < duration * fps:
        ret, frame = cap.read()
        if not ret:
            break
        
        out.write(frame)

        frame_idx += 1
        if frame_idx % (fps*SAMPLES_FPS) != 3:
            continue
        
        keyframe_file = os.path.join(temp_dir, f"{t}_keyframe_{len(key_frames)}.jpg")
        cv2.imwrite(keyframe_file, frame)
        key_frames.append(keyframe_file)
        if len(key_frames) >= 2:
            break

    # 清理资源
    cap.release()
    out.release()
    print(f"✅ 抽取完成 {len(key_frames)} 帧: {key_frames}")
    
    return {
            "frames": key_frames, 
            "video": video_file 
            }
